/**
 * NetFlow - OpenWRT 网络流量管理系统
 * 优先级调度模块
 * 
 * 本文件实现了差分优先级调度功能，
 * 包括多维优先级计算、队列管理和带宽分配。
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <syslog.h>
#include <errno.h>
#include <time.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netlink/netlink.h>
#include <netlink/route/link.h>
#include <netlink/route/tc.h>
#include <netlink/route/qdisc.h>
#include <netlink/route/class.h>
#include <netlink/route/cls/u32.h>

#include "netflow.h"
#include "priority_scheduler.h"
#include "traffic_classifier.h"
#include "util.h"

/* 全局变量 */
static char interface[IFNAMSIZ];
static uint32_t total_bandwidth;
static int if_index;
static struct nl_sock *nl_sock;

/**
 * 初始化优先级调度器
 */
int init_priority_scheduler(const char *ifname, uint32_t bandwidth) {
    syslog(LOG_INFO, "初始化优先级调度器...");
    
    // 保存接口名称和带宽
    strncpy(interface, ifname, IFNAMSIZ - 1);
    total_bandwidth = bandwidth;
    
    // 创建netlink套接字
    nl_sock = nl_socket_alloc();
    if (!nl_sock) {
        syslog(LOG_ERR, "无法分配netlink套接字");
        return -1;
    }
    
    // 连接到netlink路由协议
    if (nl_connect(nl_sock, NETLINK_ROUTE) < 0) {
        syslog(LOG_ERR, "无法连接到netlink路由协议");
        nl_socket_free(nl_sock);
        return -1;
    }
    
    // 获取接口索引
    struct nl_cache *link_cache;
    if (rtnl_link_alloc_cache(nl_sock, AF_UNSPEC, &link_cache) < 0) {
        syslog(LOG_ERR, "无法获取链路缓存");
        nl_socket_free(nl_sock);
        return -1;
    }
    
    struct rtnl_link *link = rtnl_link_get_by_name(link_cache, interface);
    if (!link) {
        syslog(LOG_ERR, "无法找到接口: %s", interface);
        nl_cache_free(link_cache);
        nl_socket_free(nl_sock);
        return -1;
    }
    
    if_index = rtnl_link_get_ifindex(link);
    rtnl_link_put(link);
    nl_cache_free(link_cache);
    
    // 设置TC队列规则
    if (setup_tc_qdisc() < 0) {
        syslog(LOG_ERR, "无法设置TC队列规则");
        nl_socket_free(nl_sock);
        return -1;
    }
    
    syslog(LOG_INFO, "优先级调度器初始化成功");
    return 0;
}

/**
 * 清理优先级调度器
 */
void cleanup_priority_scheduler(void) {
    syslog(LOG_INFO, "清理优先级调度器...");
    
    // 删除TC队列规则
    cleanup_tc_qdisc();
    
    // 释放netlink套接字
    if (nl_sock) {
        nl_socket_free(nl_sock);
        nl_sock = NULL;
    }
    
    syslog(LOG_INFO, "优先级调度器已清理");
}

/**
 * 设置TC队列规则
 */
int setup_tc_qdisc(void) {
    // 清除现有队列规则
    cleanup_tc_qdisc();
    
    // 创建HTB根队列
    if (setup_htb_root() < 0) {
        return -1;
    }
    
    // 创建主类，设置总带宽
    if (setup_htb_main_class() < 0) {
        return -1;
    }
    
    // 创建优先级类
    if (setup_htb_priority_classes() < 0) {
        return -1;
    }
    
    // 为每个类添加SFQ队列规则
    if (setup_sfq_qdisc_for_classes() < 0) {
        return -1;
    }
    
    // 添加过滤器
    if (setup_filters() < 0) {
        return -1;
    }
    
    return 0;
}

/**
 * 清除TC队列规则
 */
void cleanup_tc_qdisc(void) {
    struct nl_cache *qdisc_cache;
    struct rtnl_qdisc *qdisc;
    
    // 获取现有队列规则
    if (rtnl_qdisc_alloc_cache(nl_sock, if_index, &qdisc_cache) < 0) {
        syslog(LOG_WARNING, "无法获取队列规则缓存");
        return;
    }
    
    // 删除根队列规则
    qdisc = rtnl_qdisc_alloc();
    if (!qdisc) {
        syslog(LOG_WARNING, "无法分配队列规则");
        nl_cache_free(qdisc_cache);
        return;
    }
    
    rtnl_tc_set_link(TC_CAST(qdisc), if_index);
    rtnl_tc_set_parent(TC_CAST(qdisc), TC_H_ROOT);
    
    if (rtnl_qdisc_delete(nl_sock, qdisc) < 0) {
        syslog(LOG_WARNING, "无法删除根队列规则");
    }
    
    rtnl_qdisc_put(qdisc);
    nl_cache_free(qdisc_cache);
}

/**
 * 设置HTB根队列
 */
int setup_htb_root(void) {
    struct rtnl_qdisc *qdisc;
    int err;
    
    qdisc = rtnl_qdisc_alloc();
    if (!qdisc) {
        syslog(LOG_ERR, "无法分配队列规则");
        return -1;
    }
    
    rtnl_tc_set_link(TC_CAST(qdisc), if_index);
    rtnl_tc_set_parent(TC_CAST(qdisc), TC_H_ROOT);
    rtnl_tc_set_handle(TC_CAST(qdisc), 0x10000);
    rtnl_qdisc_set_kind(qdisc, "htb");
    
    // 设置默认类（最低优先级）
    rtnl_htb_set_defcls(qdisc, 40);
    
    err = rtnl_qdisc_add(nl_sock, qdisc, NLM_F_CREATE | NLM_F_REPLACE);
    rtnl_qdisc_put(qdisc);
    
    if (err < 0) {
        syslog(LOG_ERR, "无法添加HTB根队列: %s", nl_geterror(err));
        return -1;
    }
    
    return 0;
}

/**
 * 设置HTB主类
 */
int setup_htb_main_class(void) {
    struct rtnl_class *cls;
    int err;
    
    cls = rtnl_class_alloc();
    if (!cls) {
        syslog(LOG_ERR, "无法分配类");
        return -1;
    }
    
    rtnl_tc_set_link(TC_CAST(cls), if_index);
    rtnl_tc_set_parent(TC_CAST(cls), 0x10000);
    rtnl_tc_set_handle(TC_CAST(cls), 0x10001);
    rtnl_class_set_kind(cls, "htb");
    
    // 设置类参数
    rtnl_htb_set_rate(cls, total_bandwidth);
    rtnl_htb_set_ceil(cls, total_bandwidth);
    
    err = rtnl_class_add(nl_sock, cls, NLM_F_CREATE | NLM_F_REPLACE);
    rtnl_class_put(cls);
    
    if (err < 0) {
        syslog(LOG_ERR, "无法添加HTB主类: %s", nl_geterror(err));
        return -1;
    }
    
    return 0;
}

/**
 * 设置HTB优先级类
 */
int setup_htb_priority_classes(void) {
    struct rtnl_class *cls;
    int err;
    
    // 优先级类带宽分配
    struct {
        uint32_t min_rate;   // 最小保证带宽（总带宽的百分比）
        uint32_t max_rate;   // 最大可用带宽（总带宽的百分比）
        uint8_t prio;        // 优先级（越小越优先）
    } prio_classes[] = {
        { 40, 100, 1 },  // 实时交互类（最高优先级）
        { 30, 80, 2 },   // 流媒体类（次高优先级）
        { 20, 60, 3 },   // 普通数据类（中等优先级）
        { 10, 40, 4 }    // 背景传输类（最低优先级）
    };
    
    for (int i = 0; i < 4; i++) {
        cls = rtnl_class_alloc();
        if (!cls) {
            syslog(LOG_ERR, "无法分配类");
            return -1;
        }
        
        rtnl_tc_set_link(TC_CAST(cls), if_index);
        rtnl_tc_set_parent(TC_CAST(cls), 0x10001);
        rtnl_tc_set_handle(TC_CAST(cls), 0x10010 + (i * 0x10));
        rtnl_class_set_kind(cls, "htb");
        
        // 设置类参数
        uint32_t min_rate = total_bandwidth * prio_classes[i].min_rate / 100;
        uint32_t max_rate = total_bandwidth * prio_classes[i].max_rate / 100;
        
        rtnl_htb_set_rate(cls, min_rate);
        rtnl_htb_set_ceil(cls, max_rate);
        rtnl_htb_set_prio(cls, prio_classes[i].prio);
        
        err = rtnl_class_add(nl_sock, cls, NLM_F_CREATE | NLM_F_REPLACE);
        rtnl_class_put(cls);
        
        if (err < 0) {
            syslog(LOG_ERR, "无法添加HTB优先级类 %d: %s", i+1, nl_geterror(err));
            return -1;
        }
    }
    
    return 0;
}

/**
 * 为每个类添加SFQ队列规则
 */
int setup_sfq_qdisc_for_classes(void) {
    struct rtnl_qdisc *qdisc;
    int err;
    
    for (int i = 0; i < 4; i++) {
        qdisc = rtnl_qdisc_alloc();
        if (!qdisc) {
            syslog(LOG_ERR, "无法分配队列规则");
            return -1;
        }
        
        rtnl_tc_set_link(TC_CAST(qdisc), if_index);
        rtnl_tc_set_parent(TC_CAST(qdisc), 0x10010 + (i * 0x10));
        rtnl_tc_set_handle(TC_CAST(qdisc), 0x20000 + (i * 0x10000));
        rtnl_qdisc_set_kind(qdisc, "sfq");
        
        // 设置SFQ参数
        rtnl_sfq_set_perturb(qdisc, 10);  // 每10秒重新哈希
        rtnl_sfq_set_quantum(qdisc, 1514);  // 最大数据包大小
        
        err = rtnl_qdisc_add(nl_sock, qdisc, NLM_F_CREATE | NLM_F_REPLACE);
        rtnl_qdisc_put(qdisc);
        
        if (err < 0) {
            syslog(LOG_ERR, "无法为优先级类 %d 添加SFQ队列规则: %s", i+1, nl_geterror(err));
            return -1;
        }
    }
    
    return 0;
}

/**
 * 设置过滤器
 */
int setup_filters(void) {
    // 根据DSCP值分类的过滤器
    struct {
        uint8_t dscp;       // DSCP值
        uint8_t mask;       // 掩码
        uint32_t classid;   // 对应的类ID
    } dscp_filters[] = {
        { 46 << 2, 0xfc, 0x10010 },  // EF (Expedited Forwarding) -> 实时交互类
        { 34 << 2, 0xfc, 0x10020 },  // AF41 -> 流媒体类
        { 18 << 2, 0xfc, 0x10030 },  // AF21 -> 普通数据类
        { 8 << 2,  0xfc, 0x10040 }   // CS1 -> 背景传输类
    };
    
    for (int i = 0; i < 4; i++) {
        if (add_dscp_filter(dscp_filters[i].dscp, dscp_filters[i].mask, dscp_filters[i].classid) < 0) {
            return -1;
        }
    }
    
    return 0;
}

/**
 * 添加DSCP过滤器
 */
int add_dscp_filter(uint8_t dscp, uint8_t mask, uint32_t classid) {
    struct rtnl_cls *cls;
    struct nl_cache *cls_cache;
    struct tc_u32_sel *sel;
    int err;
    
    cls = rtnl_cls_alloc();
    if (!cls) {
        syslog(LOG_ERR, "无法分配过滤器");
        return -1;
    }
    
    rtnl_tc_set_link(TC_CAST(cls), if_index);
    rtnl_tc_set_parent(TC_CAST(cls), 0x10000);
    rtnl_tc_set_handle(TC_CAST(cls), 0);
    rtnl_cls_set_protocol(cls, ETH_P_IP);
    rtnl_cls_set_kind(cls, "u32");
    rtnl_cls_set_prio(cls, 1);
    
    // 设置过滤器选择器
    sel = rtnl_u32_get_selector(cls);
    tc_u32_add(sel, &dscp, &mask, 1, 0, 0);
    tc_u32_add_key_tos(sel, dscp, mask);
    
    // 设置动作
    rtnl_cls_set_classid(cls, classid);
    
    err = rtnl_cls_add(nl_sock, cls, NLM_F_CREATE | NLM_F_REPLACE);
    rtnl_cls_put(cls);
    
    if (err < 0) {
        syslog(LOG_ERR, "无法添加DSCP过滤器: %s", nl_geterror(err));
        return -1;
    }
    
    return 0;
}

/**
 * 获取应用优先级
 */
int get_app_priority(int app_class) {
    switch (app_class) {
        case APP_CLASS_VIDEO_CONFERENCE:
        case APP_CLASS_VOIP:
        case APP_CLASS_ONLINE_GAME:
            return PRIORITY_HIGHEST;  // 实时交互类
            
        case APP_CLASS_VIDEO_STREAMING:
        case APP_CLASS_AUDIO_STREAMING:
            return PRIORITY_HIGH;     // 流媒体类
            
        case APP_CLASS_WEB_BROWSING:
        case APP_CLASS_EMAIL:
        case APP_CLASS_SSH:
            return PRIORITY_NORMAL;   // 普通数据类
            
        case APP_CLASS_FILE_DOWNLOAD:
        case APP_CLASS_P2P:
        case APP_CLASS_SYSTEM_UPDATE:
            return PRIORITY_LOW;      // 背景传输类
            
        case APP_CLASS_UNKNOWN:
        default:
            return PRIORITY_NORMAL;   // 默认为普通优先级
    }
}

/**
 * 计算多维优先级
 */
int calculate_priority(flow_info_t *flow, user_info_t *user, network_info_t *network, time_info_t *time) {
    // 基础优先级（基于应用类型）
    int base_priority = get_app_priority(flow->app_class);
    
    // 用户权重因子
    float user_factor = calculate_user_factor(user);
    
    // 网络状况因子
    float network_factor = calculate_network_factor(network, flow);
    
    // 时间因子
    float time_factor = calculate_time_factor(time, flow->app_class);
    
    // 流量特性因子
    float flow_factor = calculate_flow_factor(flow);
    
    // 计算加权优先级分数
    float priority_score = base_priority * user_factor * network_factor * time_factor * flow_factor;
    
    // 映射到优先级级别
    int priority_level = map_score_to_priority_level(priority_score);
    
    return priority_level;
}

/**
 * 计算用户因子
 */
float calculate_user_factor(user_info_t *user) {
    if (!user) return 1.0;
    
    // 基于用户优先级和设备重要性
    float user_priority = get_user_priority(user->user_id);
    float device_importance = get_device_importance(user->device_id);
    
    return 0.7 * user_priority + 0.3 * device_importance;
}

/**
 * 获取用户优先级
 */
float get_user_priority(int user_id) {
    // 在实际应用中，应从用户配置或数据库获取
    // 这里使用简单的映射关系
    switch (user_id) {
        case 1:  return 1.5;  // 高优先级用户
        case 2:  return 1.2;  // 中高优先级用户
        case 3:  return 1.0;  // 标准优先级用户
        case 4:  return 0.8;  // 低优先级用户
        default: return 1.0;  // 默认为标准优先级
    }
}

/**
 * 获取设备重要性
 */
float get_device_importance(int device_id) {
    // 在实际应用中，应从设备配置或数据库获取
    // 这里使用简单的映射关系
    switch (device_id) {
        case 1:  return 1.5;  // 工作电脑
        case 2:  return 1.3;  // 个人电脑
        case 3:  return 1.2;  // 智能电视
        case 4:  return 1.0;  // 智能手机
        case 5:  return 0.8;  // 智能音箱
        default: return 1.0;  // 默认为标准重要性
    }
}

/**
 * 计算网络因子
 */
float calculate_network_factor(network_info_t *network, flow_info_t *flow) {
    if (!network) return 1.0;
    
    // 基于网络负载和拥塞状况
    float load_factor = 1.0 - (network->current_load / network->capacity);
    float congestion_factor = 1.0 - network->congestion_level;
    
    // 不同应用对网络状况的敏感度不同
    float sensitivity = get_application_sensitivity(flow->app_class);
    
    return load_factor * (1.0 - sensitivity) + congestion_factor * sensitivity;
}

/**
 * 获取应用敏感度
 */
float get_application_sensitivity(int app_class) {
    switch (app_class) {
        case APP_CLASS_VIDEO_CONFERENCE:
        case APP_CLASS_VOIP:
        case APP_CLASS_ONLINE_GAME:
            return 0.9;  // 高度敏感
            
        case APP_CLASS_VIDEO_STREAMING:
        case APP_CLASS_AUDIO_STREAMING:
            return 0.7;  // 较敏感
            
        case APP_CLASS_WEB_BROWSING:
        case APP_CLASS_EMAIL:
        case APP_CLASS_SSH:
            return 0.5;  // 中等敏感
            
        case APP_CLASS_FILE_DOWNLOAD:
        case APP_CLASS_P2P:
        case APP_CLASS_SYSTEM_UPDATE:
            return 0.2;  // 不太敏感
            
        case APP_CLASS_UNKNOWN:
        default:
            return 0.5;  // 默认为中等敏感
    }
}

/**
 * 计算时间因子
 */
float calculate_time_factor(time_info_t *time, int app_class) {
    if (!time) return 1.0;
    
    // 基于当前时间段的重要性
    float time_importance = get_time_importance(time->hour, app_class);
    
    // 工作时间 vs 非工作时间
    float work_hour_factor = is_work_hour(time->hour, time->day_of_week) ? 1.2 : 0.8;
    
    return time_importance * work_hour_factor;
}

/**
 * 获取时间重要性
 */
float get_time_importance(int hour, int app_class) {
    // 工作时间（9:00-18:00）
    if (hour >= 9 && hour < 18) {
        switch (app_class) {
            case APP_CLASS_VIDEO_CONFERENCE:
            case APP_CLASS_VOIP:
            case APP_CLASS_SSH:
                return 1.5;  // 工作时间内工作应用重要性高
                
            case APP_CLASS_VIDEO_STREAMING:
            case APP_CLASS_AUDIO_STREAMING:
            case APP_CLASS_ONLINE_GAME:
                return 0.8;  // 工作时间内娱乐应用重要性低
                
            default:
                return 1.0;
        }
    } 
    // 晚间（18:00-23:00）
    else if (hour >= 18 && hour < 23) {
        switch (app_class) {
            case APP_CLASS_VIDEO_STREAMING:
            case APP_CLASS_AUDIO_STREAMING:
            case APP_CLASS_ONLINE_GAME:
                return 1.3;  // 晚间娱乐应用重要性高
                
            case APP_CLASS_VIDEO_CONFERENCE:
            case APP_CLASS_VOIP:
                return 1.2;  // 晚间通讯应用仍较重要
                
            default:
                return 1.0;
        }
    }
    // 夜间/凌晨（23:00-9:00）
    else {
        return 1.0;  // 夜间所有应用同等重要
    }
}

/**
 * 是否为工作时间
 */
int is_work_hour(int hour, int day_of_week) {
    // 工作日（周一至周五）
    if (day_of_week >= 1 && day_of_week <= 5) {
        // 工作时间（9:00-18:00）
        if (hour >= 9 && hour < 18) {
            return 1;
        }
    }
    
    return 0;
}

/**
 * 计算流量特性因子
 */
float calculate_flow_factor(flow_info_t *flow) {
    if (!flow) return 1.0;
    
    // 基于流量的持续时间、数据量和突发性
    float duration_factor = calculate_duration_factor(flow->stats.last_seen - flow->stats.first_seen);
    float data_volume_factor = calculate_data_volume_factor(flow->stats.byte_count);
    float burst_factor = calculate_burst_factor(flow);
    
    return 0.4 * duration_factor + 0.3 * data_volume_factor + 0.3 * burst_factor;
}

/**
 * 计算持续时间因子
 */
float calculate_duration_factor(uint64_t duration) {
    // 短时流（<1秒）
    if (duration < 1000) {
        return 1.2;  // 短时流可能是交互性应用，给予更高优先级
    }
    // 中等时长流（1秒-5分钟）
    else if (duration < 300000) {
        return 1.0;  // 标准优先级
    }
    // 长时流（>5分钟）
    else {
        return 0.8;  // 长时流可能是后台传输，给予较低优先级
    }
}

/**
 * 计算数据量因子
 */
float calculate_data_volume_factor(uint64_t byte_count) {
    // 小流量（<100KB）
    if (byte_count < 100 * 1024) {
        return 1.2;  // 小流量可能是控制消息，给予更高优先级
    }
    // 中等流量（100KB-10MB）
    else if (byte_count < 10 * 1024 * 1024) {
        return 1.0;  // 标准优先级
    }
    // 大流量（>10MB）
    else {
        return 0.8;  // 大流量可能是文件下载，给予较低优先级
    }
}

/**
 * 计算突发因子
 */
float calculate_burst_factor(flow_info_t *flow) {
    // 检查是否有足够的样本
    if (flow->stats.iat_sample_count < 5) {
        return 1.0;  // 样本不足，返回标准值
    }
    
    // 计算包间隔时间的变异系数（标准差/均值）
    double sum = 0, sum_sq = 0;
    
    for (int i = 0; i < flow->stats.iat_sample_count; i++) {
        sum += flow->stats.packet_inter_arrival_times[i];
        sum_sq += (double)flow->stats.packet_inter_arrival_times[i] * flow->stats.packet_inter_arrival_times[i];
    }
    
    double mean = sum / flow->stats.iat_sample_count;
    double variance = (sum_sq / flow->stats.iat_sample_count) - (mean * mean);
    double std_dev = sqrt(variance);
    double cv = std_dev / mean;
    
    // 低变异（稳定流）
    if (cv < 0.5) {
        return 1.0;  // 标准优先级
    }
    // 中等变异
    else if (cv < 1.0) {
        return 1.1;  // 略高优先级
    }
    // 高变异（突发流）
    else {
        return 1.2;  // 较高优先级
    }
}

/**
 * 映射分数到优先级级别
 */
int map_score_to_priority_level(float priority_score) {
    if (priority_score >= 1.5) {
        return PRIORITY_HIGHEST;
    } else if (priority_score >= 1.2) {
        return PRIORITY_HIGH;
    } else if (priority_score >= 0.8) {
        return PRIORITY_NORMAL;
    } else {
        return PRIORITY_LOW;
    }
}

/**
 * 动态调整带宽分配
 */
void dynamic_bandwidth_allocation(float load_factor) {
    // 根据负载因子调整带宽分配
    uint32_t min_bw[4], max_bw[4];
    
    if (load_factor < 0.5) {
        // 负载较低，放宽带宽限制
        min_bw[0] = total_bandwidth * 30 / 100;  // 实时交互类
        min_bw[1] = total_bandwidth * 25 / 100;  // 流媒体类
        min_bw[2] = total_bandwidth * 25 / 100;  // 普通数据类
        min_bw[3] = total_bandwidth